/*******************************************************************************
 * This file is part of Crunch Network.
 *
 * Crunch Network is free software: you can redistribute it and/or modify it under the
 * terms of the GNU Lesser General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later version.
 *
 * Crunch Network is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
 * PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License along with
 * Crunch Network.  If not, see <http://www.gnu.org/licenses/>.
 ******************************************************************************/

package com.crunch.network.util;

import java.nio.BufferOverflowException;
import java.nio.ByteBuffer;

/**
 * Used to write to buffers with single bit granularity.
 */
public class BitStreamWriter {
	/**
	 * Creates a bit stream writer which will write to the specified buffer.
	 *
	 * @param   buffer the buffer to read to.
	 * @param   bitPosition the position of the first writable bit in the buffer.
	 * @param   bitLimit the position of the first unwritable bit in the buffer.
	 */
	public BitStreamWriter(ByteBuffer buffer, int bitPosition, int bitLimit) {
		this.buffer = buffer;

		int bytePosition = bitPosition / Byte.SIZE;
		buffer.position(bytePosition);
		bitOffset = bitPosition % Byte.SIZE;

		int byteLimit = (bitLimit + Byte.SIZE - 1) / Byte.SIZE;
		buffer.limit(byteLimit);
		unwritableEndBits = (byteLimit * Byte.SIZE) - bitLimit;
	}

	/**
	 * Writes {@code bitCount} bits to the stream and increments the position by {@code bitCount}.
	 *
	 * @param   value the value to write to the stream.
	 * @param   bitCount the number of bits to write; must be between 1 and 32.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeBits(int value, int bitCount) {
		if (bitCount <= 0 || bitCount > 32) {
			throw new IllegalArgumentException("Invalid bit count");
		}

		if (remaining() < bitCount) {
			throw new BufferOverflowException();
		}

		// write bits with MSB first
		int bitsRemaining = bitCount;
		// write bits in chunks
		while (bitsRemaining > 0) {
			// determine how many bits to write next
			// either write the remaining bits in the whole bitfield, or the remaining bits in the byte
			int bc = Math.min(bitsRemaining, Byte.SIZE - bitOffset);
			// read the current value in the stream
			byte streamVal = buffer.get(buffer.position());

			// shift and OR to obtain the bits we want to write
			bitsRemaining -= bc;
			int valBits = (value >>> bitsRemaining) & ~(0xffffffff << bc);

			// combine with the current stream value and write back to the stream
			streamVal = (byte) (streamVal | (valBits << bitOffset));
			buffer.put(buffer.position(), streamVal);

			// increment into the stream
			bitOffset += bc;
			if (bitOffset == Byte.SIZE) {
				// we've reached the end of the byte - increment to next byte
				bitOffset = 0;
				buffer.get();
			}
		}
	}

	/**
	 * Writes a boolean value to the stream and increments the position by 1.
	 *
	 * @param   value the boolean value to write to the stream.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeBoolean(boolean value) {
		writeBits(value ? 1 : 0, 1);
	}

	/**
	 * Writes an unsigned integer value of size {@code bitCount} to the stream and increments the position by
	 * {@code bitCount}. {@code value} must not be negative.
	 *
	 * @param   value the positive integer value to write to the stream.
	 * @param   bitCount the number of bits to write; must be between 1 and 32.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeUnsignedInt(long value, int bitCount) {
		if (value < 0) {
			throw new IllegalArgumentException("Unsigned integer cannot be negative");
		}
		writeBits((int) value, bitCount);
	}

	/**
	 * Writes a signed integer value of size {@code bitCount} to the stream and increments the position by
	 * {@code bitCount}. The resulting value is sign-extended.
	 *
	 * @param   value the signed integer value to write to the stream.
	 * @param   bitCount the number of bits to write; must be between 1 and 32.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeInt(int value, int bitCount) {
		writeBits(value, bitCount);
	}

	/**
	 * Writes a float value to the stream and increments the position by {@link Float#SIZE}.
	 *
	 * @param   value the float value to write to the stream.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeFloat(float value) {
		int bits = Float.floatToIntBits(value);
		writeInt(bits, Float.SIZE);
	}

	/**
	 * Writes a quantized unsigned normalized float value of size {@code bitCount} to the stream and increments the
	 * position by {@code bitCount}.
	 *
	 * @param   value the unsigned normalized float value to write to the stream. The value is automatically clamped to
	 *          the range {@code [0.0f, 1.0f]}.
	 * @param   bitCount the number of bits to write; must be between 1 and 32.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeUNormFloat(float value, int bitCount) {
		writeUnsignedInt(Compressor.compressUNormFloat(value, bitCount) & 0x00000000ffffffffL, bitCount);
	}

	/**
	 * Writes a quantized signed normalized float value of size {@code bitCount} to the stream and increments the
	 * position by {@code bitCount}.
	 *
	 * @param   value the signed normalized float value to write to the stream. The value is automatically clamped to
	 *          the range {@code [-1.0f, 1.0f]}.
	 * @param   bitCount the number of bits to write; must be between 1 and 32.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeSNormFloat(float value, int bitCount) {
		writeUnsignedInt(Compressor.compressUNormFloat(value, bitCount) & 0x00000000ffffffffL, bitCount);
	}

	/**
	 * Writes a quantized wrapped normalized float value of size {@code bitCount} to the stream and increments the
	 * position by {@code bitCount}. This method differs from {@link BitStreamWriter#writeUNormFloat(float, int)} in
	 * that the domain of this method is a half-open interval.
	 *
	 * @param   value the wrapped normalized float value to write to the stream. The value is automatically wrapped to
	 *          the range {@code [0.0f, 1.0f)}.
	 * @param   bitCount the number of bits to write; must be between 1 and 32.
	 * @throws  BufferOverflowException if the limit would be exceeded with the write.
	 */
	public void writeWrappedNormFloat(float value, int bitCount) {
		writeUnsignedInt(Compressor.compressWrappedNormFloat(value, bitCount) & 0x00000000ffffffffL, bitCount);
	}

	/**
	 * Returns the limit in bits.
	 *
	 * @return  the limit in bits.
	 */
	public int limit() {
		return (buffer.limit() * Byte.SIZE) - unwritableEndBits;
	}

	/**
	 * Returns the current position in bits.
	 *
	 * @return  the current position in bits.
	 */
	public int position() {
		return (buffer.position() * Byte.SIZE) + bitOffset;
	}

	/**
	 * Returns the number of bits remaining that can be written.
	 *
	 * @return  the number of bits remaining.
	 */
	public int remaining() {
		return limit() - position();
	}

	private ByteBuffer buffer;      // underlying buffer
	private int bitOffset;          // offset into the current byte
	private int unwritableEndBits;  // number of unwritable bits in the last byte of "remaining"
}
